#!/usr/bin/env sage -python

from sage.all import *
import random
import math
import numpy as np
import pylab as pl
import time as t
from scipy import stats as s

#import antigravity

start = t.time()

###Creating transition matrices###
def create_trans_matrix(L_size,gamma,q):
	r=range(0,L_size+1)
	P = np.matrix(np.zeros((L_size+1,L_size+1)))
	
	for n in range(0,L_size+1):
		if n==L_size:
			P[n,n] = 1
		else:
			P[n,n] = 1/(q**(L_size-r[n]))
			P[n,n+1] = 1-1/(q**(L_size-r[n]))
			#P[n,n] = (1-gamma)*(1/(q**(L_size-r[n])))
			#P[n,n+1] = float(1)-(1-gamma)*(1/(q**(L_size-r[n]-1)))
	return P


def Pm(m,n,r,q):
	S = vector(ZZ,n+1)
	S[0]=1
	matrix=create_trans_matrix(n,1,q)
	pmf = np.array(S * (matrix**m))
	#print matrix
	#print pmf[0,:]
	return pmf[0,r]


transmissions=185

q_liste=[float(2),float(256)]
qs = 2



gamma_liste=[[0.3,0.7],[0.4,0.6],[0.5,0.5]]
gammas = 3

L1_size=32
L2_size=64



S1 = vector(ZZ,L1_size+1)
S1[0]=1
S2 = vector(ZZ,L2_size+1)
S2[0]=1

for j in range(0,gammas):
	gamma = gamma_liste[j]
	
	
	for a in range(0,qs):
		q = q_liste[a]
		P1 = create_trans_matrix(L1_size,1,q)
		P2 = create_trans_matrix(L2_size,1,q)
		result=np.zeros(transmissions+1)
		result1=np.zeros(transmissions+1)
		
		for N in range(0,transmissions+1):
			print N
			for n in range(0,N+1):
				pmf1 = np.array(S1 * (P1**n))
				pmf2 = np.array(S2 * (P2**(N-n)))
			
				layer1_p = s.binom.pmf(n, N, gamma[0]) * pmf1[0,L1_size]
			
				result[N] += s.binom.pmf(n, N, gamma[0]) * pmf1[0,L1_size] * pmf2[0,L2_size]
				result1[N] += layer1_p
	
		gam = gamma[0]*10
		path = 'saved_data_new/new_anal_L2_g%d' % gam + 'f_%d' % q + '.txt'
		L1_file = open(path,'w')
		for i in result:
			L1_file.write(str(i) +'\n')
		L1_file.close()
		
		path = 'saved_data_new/new_anal_L1_g%d' % gam + 'f_%d' % q + '.txt'
		L2_file = open(path,'w')
		for i in result1:
			L2_file.write(str(i)+'\n')
		L2_file.close()
		t = np.arange(0, transmissions+1)
		pl.plot(t,result)
		pl.plot(t,result1)
		pl.grid(True)
		pl.xticks(np.arange(0,transmissions+1,10))
		pl.savefig('simple_plot%d' % j)


print "This script took ", (t.time()-start)/60, " minutes."